#include "ollama_client.hpp"
#include <curl/curl.h>
#include <nlohmann/json.hpp>
#include <stdexcept>
#include <iostream>

using json = nlohmann::json;

// libcurl 写数据回调
static size_t writeCallback(void* contents, size_t size, size_t nmemb, void* userp) {
    auto* s = static_cast<std::string*>(userp);
    size_t total = size * nmemb;
    s->append(static_cast<char*>(contents), total);
    return total;
}

namespace ollama {

std::string chat_with_ollama(const std::string& user_content,
                             const std::string& system_prompt,
                             const std::string& model_name,
                             const std::string& host)
{
    // 默认系统提示词
    const std::string default_system =
        "";

    const std::string& system = system_prompt.empty() ? default_system : system_prompt;

    // 构造 JSON 请求体
    json payload = {
        {"model", model_name},
        {"stream", false},
        {"messages", json::array({
            {{"role","system"},{"content", system}},
            {{"role","user"},  {"content", user_content}}
        })}
    };
    std::string post_data = payload.dump();

    // libcurl 初始化
    CURL* curl = curl_easy_init();
    if (!curl) throw std::runtime_error("curl_easy_init failed");

    std::string read_buffer;
    std::string url = host + "/api/chat";

    curl_easy_setopt(curl, CURLOPT_URL, url.c_str());
    curl_easy_setopt(curl, CURLOPT_POSTFIELDS, post_data.c_str());
    curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, writeCallback);
    curl_easy_setopt(curl, CURLOPT_WRITEDATA, &read_buffer);

    struct curl_slist* headers = nullptr;
    headers = curl_slist_append(headers, "Content-Type: application/json");
    curl_easy_setopt(curl, CURLOPT_HTTPHEADER, headers);

    CURLcode res = curl_easy_perform(curl);
    curl_slist_free_all(headers);
    curl_easy_cleanup(curl);

    if (res != CURLE_OK)
        throw std::runtime_error("curl_easy_perform() failed: " +
                                 std::string(curl_easy_strerror(res)));

    // 解析返回的 JSON
    json resp = json::parse(read_buffer);
    if (!resp.contains("message") || !resp["message"].contains("content"))
        throw std::runtime_error("Invalid response format");

    return resp["message"]["content"].get<std::string>();
}

} // namespace ollama